from sklearn.datasets import fetch_openml
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from sklearn.preprocessing import StandardScaler
from skimage import io, color, transform
import matplotlib.pyplot as plt

# 加载 MNIST 数据集
mnist = fetch_openml('mnist_784', version=1, as_frame=False)
images = mnist.data
labels = mnist.target.astype(int)

# 对数据进行标准化
scaler = StandardScaler()
images = scaler.fit_transform(images)

# 划分数据集为 3:1 的训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.25, random_state=42)

# 初始化逻辑回归模型，增加 max_iter
LR_classifier = LogisticRegression(C=0.01, penalty='l2', tol=0.01, solver='lbfgs', max_iter=5000)

# 训练模型
LR_classifier.fit(X_train, y_train)

# 在测试集上进行预测
y_pred = LR_classifier.predict(X_test)

# 显示分类报告
print("Performance Report:\n", classification_report(y_test, y_pred))


# 定义函数，加载自定义图片并预测
def predict_custom_image(image_path):
    # 加载图片
    digit_img = io.imread(image_path)
    # 转为灰度图
    # digit_img = color.rgb2gray(color.rgba2rgb(digit_img))
    digit_img = color.rgb2gray(digit_img)
    # 调整大小为 28x28
    digit_img = transform.resize(digit_img, (28, 28), mode="reflect")
    # 展平为 1D 数组
    digit_flattened = digit_img.flatten().reshape(1, -1)
    # 标准化
    digit_flattened = scaler.transform(digit_flattened)

    # 预测图片
    prediction = LR_classifier.predict(digit_flattened)
    probabilities = LR_classifier.predict_proba(digit_flattened)

    # 显示图片和预测结果
    plt.imshow(digit_img, cmap='gray')
    plt.title(f"Predicted: {prediction[0]}")
    plt.axis('off')
    plt.show()

    print(f"Predicted class: {prediction[0]}")
    print(f"Prediction probabilities: {probabilities}")


# 测试自定义图片
predict_custom_image('images/digit_6.png')
